import tensorflow as tf
from tensorflow import keras

import numpy as np
import matplotlib.pyplot as plt

class_names = ['T-shirt/top', 'Trouser', 'Pullover', 'Dress', 'Coat',
               'Sandal', 'Shirt', 'Sneaker', 'Bag', 'Ankle boot']

fashion_mnist = tf.keras.datasets.fashion_mnist
(train_images, train_labels), (test_images, test_labels) = fashion_mnist.load_data()

train_images = train_images / 255.0
test_images = test_images / 255.0

# 重新创建完全相同的模型，包括其权重和优化程序
model = keras.models.load_model('model/02.h5')
# 显示网络结构
model.summary()

probability_model = tf.keras.Sequential([model,
                                         tf.keras.layers.Softmax()])


def plot_value_array(i, predictions_array, true_label):
    true_label = true_label[i]
    plt.grid(False)
    plt.xticks(range(10))
    plt.yticks([])
    thisplot = plt.bar(range(10), predictions_array, color="#777777")
    plt.ylim([0, 1])
    predicted_label = np.argmax(predictions_array)

    thisplot[predicted_label].set_color('red')
    thisplot[true_label].set_color('blue')


# 6. 使用训练好的模型
# 最后，使用训练好的模型对单个图像进行预测。
# Grab an image from the test dataset.
img = test_images[1]
plt.figure()
plt.imshow(img)
plt.colorbar()
plt.grid(False)
plt.show()
print(img.shape)
# tf.keras 模型经过了优化，可同时对一个批或一组样本进行预测。因此，即便您只使用一个图像，您也需要将其添加到列表中：
# Add the image to a batch where it's the only member.
img = (np.expand_dims(img, 0))
print(img.shape)
# 现在预测这个图像的正确标签：
predictions_single = probability_model.predict(img)
print(predictions_single)

plot_value_array(1, predictions_single[0], test_labels)
_ = plt.xticks(range(10), class_names, rotation=45)
plt.show()

# keras.Model.predict 会返回一组列表，每个列表对应一批数据中的每个图像
print("在批次中获取对我们（唯一）图像的预测", np.argmax(predictions_single[0]), " - ",
      class_names[np.argmax(predictions_single[0])])
